package it.uniroma2.util.math;


/*************************************************************************
 *  Compilation:  javac FFT.java
 *  Execution:    java FFT N
 *  Dependencies: Complex.java
 *  
 *  http://introcs.cs.princeton.edu/java/97data/FFT.java.html
 *  
 *  Compute the FFT and inverse FFT of a length N complex sequence.
 *  Bare bones implementation that runs in O(N log N) time. Our goal
 *  is to optimize the clarity of the code, rather than performance.
 *
 *  Limitations
 *  -----------
 *   -  assumes N is a power of 2
 *
 *   -  not the most memory efficient algorithm (because it uses
 *      an object type for representing complex numbers and because
 *      it re-allocates memory for the subarray, instead of doing
 *      in-place or reusing a single temporary array)
 *  
 *************************************************************************/

public class FFT {

	private Complex[][] wks = null;
	
    public void initialize(int grade) {
    	int dimension = (int) Math.pow(2, grade);
    	wks = new Complex [grade+1] [dimension];
    	for (int g = 0; g < grade+1; g++) {
    		dimension = (int) Math.pow(2, g);
    		for (int k = 0; k < dimension ; k++ ) {
        		double kth = -2 * k * Math.PI / dimension;
                wks[g][k] = new Complex(Math.cos(kth), Math.sin(kth));
    		}
    	}
    }
    
    public Complex[] fft(Complex[] x , int grade) {
        int N = x.length;

        // base case
        if (N == 1) return new Complex[] { x[0] };

        // radix 2 Cooley-Tukey FFT
        if (N % 2 != 0) { throw new RuntimeException("N is not a power of 2"); }

        // fft of even terms
        Complex[] even = new Complex[N/2];
        for (int k = 0; k < N/2; k++) {
            even[k] = x[2*k];
        }
        Complex[] q = fft(even,grade-1);

        // fft of odd terms
        Complex[] odd  = even;  // reuse the array
        for (int k = 0; k < N/2; k++) {
            odd[k] = x[2*k + 1];
        }
        Complex[] r = fft(odd,grade-1);

        // combine
        Complex[] y = new Complex[N];
        for (int k = 0; k < N/2; k++) {
            y[k]       = q[k].plus(wks[grade][k].times(r[k]));
            y[k + N/2] = q[k].minus(wks[grade][k].times(r[k]));
        }
        return y;
    }
    
    public Complex[] real_fft(Complex[] x , int grade) {
        int N = x.length;

        // base case
        if (N == 1) return new Complex[] { x[0] };

        // radix 2 Cooley-Tukey FFT
        if (N % 2 != 0) { throw new RuntimeException("N is not a power of 2"); }

        // fft of even terms
        Complex[] even = new Complex[N/2];
        for (int k = 0; k < N/2; k++) {
            even[k] = x[2*k];
        }
        Complex[] q = fft(even,grade-1);

        // fft of odd terms
        Complex[] odd  = even;  // reuse the array
        for (int k = 0; k < N/2; k++) {
            odd[k] = x[2*k + 1];
        }
        Complex[] r = fft(odd,grade-1);

        // combine
        Complex[] y = new Complex[N];
        y[0]       = q[0].plus(wks[grade][0].times(r[0]));
        y[0 + N/2] = q[0].minus(wks[grade][0].times(r[0]));
        for (int k = 1; k < N/2; k++) {
            y[k] = q[k].plus(wks[grade][k].times(r[k]));
            y[N-k] = y[k].conjugate();
        }
        return y;
    }

    public static boolean equal(Complex c1, Complex c2) {
    	return (c1.re() == c2.re()) && (c1.im() == c2.im());
    }
    
    public static int grado(int N) {
    	int grado = 0;
    	
    	while (N!=1) {
    		grado ++;
    		N = N/2;
    	}
    	
    	return grado;
    }
    
    // compute the inverse FFT of x[], assuming its length is a power of 2
    public Complex[] ifft(Complex[] x,int grade) {
        int N = x.length;
        Complex[] y = new Complex[N];

        // take conjugate
        for (int i = 0; i < N; i++) {
            y[i] = x[i].conjugate();
        }

        // compute forward FFT
        y = fft(y,grade);

        // take conjugate again
        for (int i = 0; i < N; i++) {
            y[i] = y[i].conjugate();
        }

        // divide by N
        for (int i = 0; i < N; i++) {
            y[i] = y[i].times(1.0 / N);
        }

        return y;

    }
    
    // compute the circular convolution of x and y
    public Complex[] cconvolve(Complex[] x, Complex[] y, int grade) {

        // should probably pad x and y with 0s so that they have same length
        // and are powers of 2
        if (x.length != y.length) { throw new RuntimeException("Dimensions don't agree"); }

        int N = x.length;

        // compute FFT of each sequence
        Complex[] a = fft(x,grade);
        Complex[] b = fft(y,grade);

        // point-wise multiply
        Complex[] c = new Complex[N];
        for (int i = 0; i < N; i++) {
            c[i] = a[i].times(b[i]);
        }

        // compute inverse FFT
        return ifft(c,grade);
    }
    
    // compute the circular convolution of x and y with only real values
    public Complex[] real_cconvolve(Complex[] x, Complex[] y, int grade) {

        // should probably pad x and y with 0s so that they have same length
        // and are powers of 2
        if (x.length != y.length) { throw new RuntimeException("Dimensions don't agree"); }

        int N = x.length;

        // compute FFT of each sequence
        Complex[] a = real_fft(x,grade);
        Complex[] b = real_fft(y,grade);

        // point-wise multiply
        Complex[] c = new Complex[N];
        for (int i = 0; i < N; i++) {
            c[i] = a[i].times(b[i]);
        }

        // compute inverse FFT
        return ifft(c,grade);
    }
}
